import os
import yaml
import pandas as pd
from autorag.utils.util import load_yaml_config, load_summary_file
from autorag.deploy import summary_df_to_yaml


def extract_best_config_with_encoding(trial_path, output_path=None):
    """
    带编码处理的 extract_best_config 包装函数
    """
    # 读取 summary.csv
    summary_path = os.path.join(trial_path, "summary.csv")
    if not os.path.exists(summary_path):
        raise ValueError(f"summary.csv does not exist in {trial_path}.")

    trial_summary_df = load_summary_file(
        summary_path, dict_columns=["best_module_params"]
    )

    # 使用 load_yaml_config 正确读取配置文件(已包含 UTF-8 编码)
    config_yaml_path = os.path.join(trial_path, "config.yaml")
    config_dict = load_yaml_config(config_yaml_path)

    # 生成最佳配置
    yaml_dict = summary_df_to_yaml(trial_summary_df, config_dict)

    # 提取 vectordb 配置
    project_dir = os.path.dirname(trial_path)
    vectordb_config_path = os.path.join(project_dir, "resources", "vectordb.yaml")
    if os.path.exists(vectordb_config_path):
        vectordb_dict = load_yaml_config(vectordb_config_path)
        yaml_dict["vectordb"] = vectordb_dict.get("vectordb", [])
    else:
        # 使用默认配置
        yaml_dict["vectordb"] = [{
            "name": "default",
            "db_type": "chroma",
            "client_type": "persistent",
            "embedding_model": "openai",
            "collection_name": "openai",
            "path": os.path.join(project_dir, "resources", "chroma"),
        }]

        # 保存到文件
    if output_path is not None:
        with open(output_path, "w", encoding="utf-8") as f:
            yaml.safe_dump(yaml_dict, f)

    return yaml_dict


# 使用包装函数
extract_best_config_with_encoding(
    trial_path=r'D:\laern_langchain\bjq1021\content\project_dir\16',
    output_path=r'D:\laern_langchain\bjq1021\content\project_dir\16/best.yaml'
)

from autorag.deploy import extract_best_config
extract_best_config(trial_path=r'D:\laern_langchain\bjq1021\content\project_dir\16', output_path=r'D:\laern_langchain\bjq1021\content\project_dir\16/best.yaml')